import json
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from base64 import b64encode, b64decode

class EncryptionUtils:
    def __init__(self):
        self.key = b'0123456789abcdef'  # 16字节的密钥
        self.iv = b'0123456789abcdef'  # 16字节的固定IV

    def aes_encrypt(self, data):
        cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
        ct_bytes = cipher.encrypt(pad(data.encode('utf-8'), AES.block_size))
        ct = b64encode(ct_bytes).decode('utf-8')
        return ct

    def aes_decrypt(self, data):
        try:
            # b64 = json.loads(data)
            # ct = b64decode(b64['ciphertext'])
            ct=b64decode(data)
            cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
            pt = unpad(cipher.decrypt(ct), AES.block_size).decode('utf-8')
            return pt
        except:
            return None

    def encrypt_data(self, data):
        data_obj = json.loads(data)

        def encrypt_item(item):
            if isinstance(item, dict):
                encrypted_item = {}
                for key, value in item.items():
                    encrypted_item[key] = encrypt_item(value)
                return encrypted_item
            elif isinstance(item, list):
                return [encrypt_item(x) for x in item]
            else:
                # 仅当 item 是字符串(str)类型时加密
                if isinstance(item, str):
                    return self.aes_encrypt(item)
                else:
                    return item  # 非字符串保持不变

        encrypted_obj = encrypt_item(data_obj)
        return json.dumps(encrypted_obj)

